#coding:utf-8
import os
import cv2
import numpy as np
import matplotlib
matplotlib.use('TkAgg')
import matplotlib.pyplot as plt

def read_file_pic(path=os.path.dirname(os.path.abspath(__file__))):
    print "Start loading the picture"
    img_list=[]
    label_list=[]
    for file in os.listdir(path):
        if (file.endswith('.jpg')):
            img=cv2.imread(os.path.join(path, file))
            img_list.append(img)
            label_list.append(file)
    print "Done"
    return img_list,label_list


def switch_recognition(img):
    '''主要采用使用开与关时的长宽比不一样，判断方法比较简单粗暴……高度达到一定要求就判定为关，返回值1为开关为开，0为关'''
    #初步选出压板开关的大概位置
    pick = np.ones(img.shape[:2])
    pick[np.where(img[:, :, 2] <= 126)] = 0
    pick[np.where(np.sum(pick, axis=1) < 7)] = 0
    img[np.where(pick == 0)] = 0

    ret, split = cv2.threshold(img, 0, 1, cv2.THRESH_BINARY)
    split = cv2.cvtColor(split, cv2.COLOR_BGR2GRAY)
    rec, contours, hierarchy = cv2.findContours(split, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_NONE)
    contours.sort(key=np.shape, reverse=True)
    idx = 0

    for i in range(len(contours)):
        poly_img = np.zeros(img.shape, dtype=np.uint8)
        cv2.drawContours(poly_img, contours, idx, [255, 255, 255], -1)
        if np.sum(poly_img) / 255 > 300:#面积要求
            h = np.where(np.sum(poly_img[:, :, 0], axis=1))[0].shape[0]
            #w = np.where(np.sum(poly_img[:, :, 0], axis=0))[0].shape[0]
            if h > 21:
                return 0
            else:
                return 1

        idx += 1


if __name__ == '__main__':
    img_list,label_list=read_file_pic(os.path.dirname(os.path.abspath(__file__))+'/pic')
    for i in range(len(img_list)):
        img=np.copy(img_list[i])
        recog=switch_recognition(img)
        print 'label:',label_list[i]
        print 'recog:',recog